# -*- coding: utf-8 -*-

# from operator import add
#
import numpy as np
import math
import sys
import os


import numpy as np
import torch
from torchmetrics import F1Score, Precision, Recall
from torchmetrics import AUC
from loguru import logger


def show_metric(y_ans, y_pred):
    """
    TODO! batched!!
    """
    y_ans = [[label[0]] for label in y_ans]  # 可能存在重叠标签？ TODO
    y_ans = torch.tensor(y_ans)
    y_pred = torch.argmax(torch.tensor(y_pred), dim=-1).unsqueeze(dim=-1)

    assert y_ans.shape == y_pred.shape
    p = Precision()(y_pred, y_ans)
    r = Recall()(y_pred, y_ans)
    f1 = F1Score()(y_pred, y_ans)
    auc = AUC(reorder=True)(y_pred, y_ans)

    fmt = "P: {}, R:{}, F1:{}, AUC:{}".format(p, r, f1, auc)
    print(fmt)
    # logger.info(fmt)

    return p, r, f1, auc


def calc_evaluation(y_ans, y_pred):

    return show_metric(y_ans, y_pred)
